import unittest

import numpy as np
import torch

from training.tools.metrics import _fast_hist, overall_pixel_accuracy


class MetricsTestCase(unittest.TestCase):

    def test_metrics(self):
        target = np.array([
            [1, 1, 1, 1, 1],
            [1, 1, 1, 1, 1],
            [1, 1, 1, 1, 1],
            [1, 1, 1, 1, 1],
            [1, 1, 1, 1, 1],
        ])
        target = torch.from_numpy(target)

        predict = np.array([
            [1, 1, 1, 1, 1],
            [1, 0, 0, 0, 1],
            [1, 0, 0, 0, 1],
            [1, 0, 0, 0, 1],
            [1, 1, 1, 1, 1],
        ])
        predict = torch.from_numpy(predict)

        print(target.shape, predict.shape)
        hist = _fast_hist(target, predict, 2)
        overall_acc = overall_pixel_accuracy(hist)
        print(overall_acc)
        self.assertEqual(0.64, overall_acc)

    def test_metrics2(self):
        target = torch.zeros((10, 10))
        predict = torch.zeros((10, 10))
        predict[:5, :10] = 255

        print(target.shape, predict.shape)
        hist = _fast_hist(target, predict, 256)
        overall_acc = overall_pixel_accuracy(hist)
        print(overall_acc)
        self.assertEqual(0.5, overall_acc)


if __name__ == '__main__':
    unittest.main()
